# random seed for batch sampling
seed: 42

# name for this experiment in the local run directory and on wandb
exp_name: ???

valid_size: 64
# the batch size for training; for FSDP, the batch size per GPU is batch_size / (grad_accumulation_steps * num_gpus)
batch_size: 4
# the batch size during evaluation and sampling, if enabled
eval_batch_size: 8

# debug mode (disables wandb, model checkpointing, etc.)
debug: false

# the port to use for FSDP
fsdp_port: null

# wandb configuration
wandb:
  enabled: true
  entity: null
  project: "dpo-toxicity-pplm"

# to create the local run directory and cache models/datasets,
#   we will try each of these directories in order; if none exist,
#   we will create the last one and use it
local_dirs:
  - .cache

# whether or not to generate samples during evaluation; disable for FSDP/TensorParallel
#   is recommended, because they are slow

# how many model samples to generate during evaluation
n_eval_model_samples: 16

# whether to eval at the very beginning of training
do_first_eval: false

# an OmegaConf resolver that returns the local run directory, calling a function in utils.py
local_run_dir: ${get_local_run_dir:${exp_name},${local_dirs}}

# the learning rate
lr: 1e-5

# number of steps to accumulate over for each batch
#   (e.g. if batch_size=4 and gradient_accumulation_steps=2, then we will
#   accumulate gradients over 2 microbatches of size 2)
gradient_accumulation_steps: 1

# the maximum gradient norm to clip to
max_grad_norm: 10.0

# the maximum allowed length for an input (prompt + response)
max_length: 256
max_new_tokens: 64

# the maximum allowed length for a prompt
max_prompt_length: 64

# the number of epochs to train for; if null, must specify n_examples
n_epochs: 5

# the trainer class to use (e.g. BasicTrainer, FSDPTrainer, TensorParallelTrainer)
trainer: BasicTrainer

# The optimizer to use; we use RMSprop because it works about as well as Adam and is more memory-efficient
optimizer: RMSprop

# number of linear warmup steps for the learning rate
warmup_steps: 150

# whether or not to use activation/gradient checkpointing
activation_checkpointing: false

# evaluate and save model every eval_every steps
eval_every: 100
save_every: 100
validation_metric: "loss/valid"
validation_direction: "min"
validation_patience: 30

sample_during_eval: false
sample_every: 2000

# prevent wandb from logging more than once per minimum_log_interval_secs
minimum_log_interval_secs: 2.0

defaults:
- _self_
- model: gpt2-medium
- loss: dpo # which loss function, either sft or dpo (specify loss.beta if using dpo)
